feat(attention): integrate Sol-Attn for Wan2.1 - #24

Description

Integrate Sol-Attn as a built-in TeleFuser attention backend and add an official-profile adaptation for Wan2.1.

The PR vendors the Sol-Attn kernels into telefuser.kernel.sol_attn, exposes them through telefuser.ops.attention, and adds Wan2.1 Morton3D token ordering, dense warm-up guards, threshold configuration, and automatic KV split selection.

Motivation

Wan2.1 video generation produces long token sequences where dense attention is expensive. Sol-Attn dynamically selects sparse attention blocks and reduces attention latency while retaining output quality.

Keeping the kernels inside TeleFuser also avoids coupling Sol-Attn to the independently packaged tf-kernel distribution. Unsupported inputs automatically use the existing dense fallback.

Type of Change

• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality not to work as expected)
• Performance improvement
• Code refactoring
• Documentation update
• Other (please describe):

Changes Made

• Vendor the official Sol-Attn implementation under telefuser/kernel/sol_attn, including license and third-party notices.
• Add AttnImplType.SOL_ATTN and AttentionConfig.sol_attention() with the official Wan2.1 defaults:
  • dense_timesteps=10
  • dense_layers=1
  • tau=1.0
  • threshold_type="diag"
  • kv_splits="auto"
• Route Sol-Attn through the public, compile-aware telefuser.ops.attention layer with native dense fallbacks.
• Adapt Wan2.1 to the official global Morton3D ordering for tokens, RoPE frequencies, and per-token modulation.
• Restore the original token order after the DiT blocks, including pipeline-parallel execution.
• Use forward denoising-step indices for the official Sol-Attn dense warm-up schedule.
• Automatically select four KV splits for long SM90 sequences.
• Add kernel, public ops, Wan model, dense-guard, fallback, and Morton round-trip tests.
• Document configuration, packaging, compatibility, and Wan2.1 usage in the English and Chinese attention guides.

Testing

• Focused unit tests pass
• Manual Wan2.1 generation performed
• H100 kernel correctness tests performed
• Performance and encoded-video quality comparison performed

Test commands:

.venv/bin/python -m pytest \
  tests/unit/pipelines/wan_video \
  tests/unit/models/test_wan_video_sol_attention.py \
  tests/unit/models/test_wan_video_vae_spatial.py \
  tests/unit/ops/test_sol_attention.py \
  tests/unit/kernel/test_sol_attn.py -q

ruff check \
  telefuser/core/config.py \
  telefuser/models/wan_video_dit.py \
  telefuser/ops/attention \
  telefuser/pipelines/wan_video \
  tests/unit/models/test_wan_video_sol_attention.py \
  tests/unit/ops/test_sol_attention.py \
  tests/unit/kernel/test_sol_attn.py

ruff format --check \
  telefuser/core/config.py \
  telefuser/models/wan_video_dit.py \
  telefuser/ops/attention/attention_impl.py \
  telefuser/ops/attention/backends.py \
  telefuser/pipelines/wan_video/wan21_video.py \
  telefuser/pipelines/wan_video/single_dit_denoising.py \
  tests/unit/models/test_wan_video_sol_attention.py \
  tests/unit/ops/test_sol_attention.py \
  tests/unit/kernel/test_sol_attn.py

git diff --check main...sol-attn

Result: 29 passed. The H100 tests execute the actual Sol-Attn kernel and compare its output with PyTorch SDPA.

Checklist

• Code follows the project's coding standards (ruff)
• Commit-time pre-commit hooks pass
• Full repository test suite passes (pytest tests/ was not run)
• New tests added for new functionality
• Documentation updated
• Commit messages are clear and descriptive
• PR title follows the convention: [Feature] Integrate Sol-Attn for Wan2.1

Related Issues

N/A

Additional Notes

• Sol-Attn is built into TeleFuser and is not added to tf-kernel.
• No TeleFuser kernel extra or external wheel dependency is introduced.
• Eligible calls require contiguous, noncausal BF16 Q/K/V tensors with equal shapes and head dimension 128.
• Dense warm-up steps, unsupported inputs, unavailable runtimes, and kernel execution failures use the existing dense attention fallback.
• The vendored source is based on the official Sol-Engine implementation at commit 8a26fb0ec9e353125ead798cb2e312d5ce48cded.
• Third-party provenance and licenses are included in THIRD_PARTY_NOTICES.md.

GPU Architecture Support

• SM80+ (Ampere and Ada through the Triton fallback; not validated in this PR)
• SM90 (Hopper H100; CuTe path validated)
• SM100+ (Blackwell CuTe paths included; not validated in this PR)

Performance Impact

Measured against main using Wan2.1-T2V-1.3B on one H100:

Encoded-video comparison against the FlashAttention 4 output:

The comparison used identical prompts, seed, scheduler parameters, output dimensions, frame count, and encoding settings. Results are from a single H100 run and may vary by environment.